diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..e721f86e08749068b60d6510055d78046ff82732 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,40 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+040000/model_state/ocdbt.process_0/d/de047e9eccf6eefb791a571eae771116 filter=lfs diff=lfs merge=lfs -text
+190000/model_state/ocdbt.process_0/d/90986d89119a1a9d0a93a58c8ffd3aaf filter=lfs diff=lfs merge=lfs -text
+180000/model_state/ocdbt.process_0/d/b8a20e22aa870a02f9ebdd120566cee1 filter=lfs diff=lfs merge=lfs -text
+100000/model_state/ocdbt.process_0/d/4908b67a51e8f6b19d57e06dbcc3a6d6 filter=lfs diff=lfs merge=lfs -text
+195000/model_state/ocdbt.process_0/d/805b4d2cbff076e49d4272cb4de1c3c2 filter=lfs diff=lfs merge=lfs -text
+120000/model_state/ocdbt.process_0/d/dca42bda6ce575cf0e3facd6ccefef2e filter=lfs diff=lfs merge=lfs -text
+190000/model_state/ocdbt.process_0/d/682083c1dd6cfab1449e6f9c0956b612 filter=lfs diff=lfs merge=lfs -text
+140000/model_state/ocdbt.process_0/d/8dd46fb95b571ed876e66163c846288e filter=lfs diff=lfs merge=lfs -text
+140000/model_state/ocdbt.process_0/d/23085a1042588716041c3d37e0492fed filter=lfs diff=lfs merge=lfs -text
+160000/model_state/ocdbt.process_0/d/9e1ef595d4e11b2fc1a9c26f5880e9cf filter=lfs diff=lfs merge=lfs -text
+195000/model_state/ocdbt.process_0/d/8d25c17882bc09ef12ec3d5d753f6a4a filter=lfs diff=lfs merge=lfs -text
+080000/model_state/ocdbt.process_0/d/2b34c992eeed9ced9ff9fc71c9299a07 filter=lfs diff=lfs merge=lfs -text
+060000/model_state/ocdbt.process_0/d/5d6955d327966b061974e60856f65de2 filter=lfs diff=lfs merge=lfs -text
+160000/model_state/ocdbt.process_0/d/1cfd62eb21313e54e9e6f1f01d1e15de filter=lfs diff=lfs merge=lfs -text
+020000/model_state/ocdbt.process_0/d/d120b3061877a9f3596f9fe1119acdfd filter=lfs diff=lfs merge=lfs -text
+080000/model_state/ocdbt.process_0/d/709cd657dccf2c5adc568db343f42725 filter=lfs diff=lfs merge=lfs -text
+060000/model_state/ocdbt.process_0/d/4ec4cbcd896ac4179022886603c5c967 filter=lfs diff=lfs merge=lfs -text
+020000/model_state/ocdbt.process_0/d/271e4e1bcfd91ec5d93dc6a910741cd4 filter=lfs diff=lfs merge=lfs -text
+200000/model_state/ocdbt.process_0/d/45b803694d1df8de743a6714a80a52e5 filter=lfs diff=lfs merge=lfs -text
+020000/model_state/ocdbt.process_0/d/dca5b28f587b225ae1457a60c4e2dd56 filter=lfs diff=lfs merge=lfs -text
+160000/model_state/ocdbt.process_0/d/e607017ce17c1af8daad4f13214b3ea6 filter=lfs diff=lfs merge=lfs -text
+200000/model_state/ocdbt.process_0/d/063057d29f7537fe5dc89d2c5438d296 filter=lfs diff=lfs merge=lfs -text
+120000/model_state/ocdbt.process_0/d/c4873dcbc98846803bc300a0644484d3 filter=lfs diff=lfs merge=lfs -text
+020000/model_state/ocdbt.process_0/d/4b7ad519a171ba543fcd45ed93281821 filter=lfs diff=lfs merge=lfs -text
+040000/model_state/ocdbt.process_0/d/2a00db4599175973782bc9ca2e6ec754 filter=lfs diff=lfs merge=lfs -text
+040000/model_state/ocdbt.process_0/d/7dab326fb2cb1f0e157e155c7c8b04cc filter=lfs diff=lfs merge=lfs -text
+120000/model_state/ocdbt.process_0/d/fa3cb93f242c469029cb4005bed01132 filter=lfs diff=lfs merge=lfs -text
+020000/model_state/ocdbt.process_0/d/d8aa5f718557c7396ade2c889d3c4294 filter=lfs diff=lfs merge=lfs -text
+060000/model_state/ocdbt.process_0/d/c81fa37784c7b44640b76c45bd180d65 filter=lfs diff=lfs merge=lfs -text
+080000/model_state/ocdbt.process_0/d/7ea6d160856219e72ad3ac6bcd824739 filter=lfs diff=lfs merge=lfs -text
+195000/model_state/ocdbt.process_0/d/9a90ed05f2a1c9ce06491436f3ef894c filter=lfs diff=lfs merge=lfs -text
+140000/model_state/ocdbt.process_0/d/9974f401c505e80ccb31e492612019d9 filter=lfs diff=lfs merge=lfs -text
+160000/model_state/ocdbt.process_0/d/9aef4849f854df2df779cfb1c0555793 filter=lfs diff=lfs merge=lfs -text
+180000/model_state/ocdbt.process_0/d/ffe4cb9a3ef849cb1dd1ad05694ceb3f filter=lfs diff=lfs merge=lfs -text
+200000/model_state/ocdbt.process_0/d/ee516710be0ffac27c9c58605cf19264 filter=lfs diff=lfs merge=lfs -text
+190000/model_state/ocdbt.process_0/d/aba66df9f95c5f517e76eec83302d73d filter=lfs diff=lfs merge=lfs -text
+100000/model_state/ocdbt.process_0/d/98105609c0838c418037c239d750afd7 filter=lfs diff=lfs merge=lfs -text
diff --git a/020000/model_state/d/fc9e9da0ba0c8adb329dd7a34cac4f2d b/020000/model_state/d/fc9e9da0ba0c8adb329dd7a34cac4f2d
new file mode 100644
index 0000000000000000000000000000000000000000..b660da2ac1dac05e6135936f5cda06f0a55de401
Binary files /dev/null and b/020000/model_state/d/fc9e9da0ba0c8adb329dd7a34cac4f2d differ
diff --git a/020000/model_state/ocdbt.process_0/d/271e4e1bcfd91ec5d93dc6a910741cd4 b/020000/model_state/ocdbt.process_0/d/271e4e1bcfd91ec5d93dc6a910741cd4
new file mode 100644
index 0000000000000000000000000000000000000000..dd29d577f060123933fc38829a2a5f991dee6c5b
--- /dev/null
+++ b/020000/model_state/ocdbt.process_0/d/271e4e1bcfd91ec5d93dc6a910741cd4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d73156207fb736299a6bc06b88491b168a4decac725aabfcef43d32e3a17b549
+size 5877760
diff --git a/020000/model_state/ocdbt.process_0/d/34d43e4701b8c2d163f8dc88606fac73 b/020000/model_state/ocdbt.process_0/d/34d43e4701b8c2d163f8dc88606fac73
new file mode 100644
index 0000000000000000000000000000000000000000..d6d49a347d6892e5cf23585db1cac88f6633a7d1
Binary files /dev/null and b/020000/model_state/ocdbt.process_0/d/34d43e4701b8c2d163f8dc88606fac73 differ
diff --git a/020000/model_state/ocdbt.process_0/d/4b7ad519a171ba543fcd45ed93281821 b/020000/model_state/ocdbt.process_0/d/4b7ad519a171ba543fcd45ed93281821
new file mode 100644
index 0000000000000000000000000000000000000000..89cc851af67ca2bf445ccfc7f71091873327ce80
--- /dev/null
+++ b/020000/model_state/ocdbt.process_0/d/4b7ad519a171ba543fcd45ed93281821
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da6bef1dfa2d58d8811b28915646f9945be4bf1a2fa2f2711e88313e47ad768d
+size 209608704
diff --git a/020000/model_state/ocdbt.process_0/d/d120b3061877a9f3596f9fe1119acdfd b/020000/model_state/ocdbt.process_0/d/d120b3061877a9f3596f9fe1119acdfd
new file mode 100644
index 0000000000000000000000000000000000000000..d64c6974d96e60882377b4e692de35bee52d4280
--- /dev/null
+++ b/020000/model_state/ocdbt.process_0/d/d120b3061877a9f3596f9fe1119acdfd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:789dab0c998e814bd7d61966d44bfdd198386c0b122f5251e44a6dcd6fdebce8
+size 981230
diff --git a/020000/model_state/ocdbt.process_0/d/d8aa5f718557c7396ade2c889d3c4294 b/020000/model_state/ocdbt.process_0/d/d8aa5f718557c7396ade2c889d3c4294
new file mode 100644
index 0000000000000000000000000000000000000000..9a1e1ab263a81e6259d6937a0a1cc09003f9216b
--- /dev/null
+++ b/020000/model_state/ocdbt.process_0/d/d8aa5f718557c7396ade2c889d3c4294
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d29ffd7d0ccc24b5b289f7c467a6263e44665970e204984448a835945651b668
+size 371400704
diff --git a/020000/model_state/ocdbt.process_0/d/dca5b28f587b225ae1457a60c4e2dd56 b/020000/model_state/ocdbt.process_0/d/dca5b28f587b225ae1457a60c4e2dd56
new file mode 100644
index 0000000000000000000000000000000000000000..15e8cd917ac5fb7e28a9f4ac9de1a5484a60efce
--- /dev/null
+++ b/020000/model_state/ocdbt.process_0/d/dca5b28f587b225ae1457a60c4e2dd56
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bceb378a7ccf6cbda76c2491f484e386fd92e10896d9cfc4b8d3152f83fc02bf
+size 20410368
diff --git a/040000/model_state/ocdbt.process_0/d/2a00db4599175973782bc9ca2e6ec754 b/040000/model_state/ocdbt.process_0/d/2a00db4599175973782bc9ca2e6ec754
new file mode 100644
index 0000000000000000000000000000000000000000..370e5a2c36a0c90b5564e8866827bd9961a3321e
--- /dev/null
+++ b/040000/model_state/ocdbt.process_0/d/2a00db4599175973782bc9ca2e6ec754
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:738aeafe512ee4adaa5c2d82febc46862200dd4ba894af21835c4bf44e9664b3
+size 286040064
diff --git a/040000/model_state/ocdbt.process_0/d/7dab326fb2cb1f0e157e155c7c8b04cc b/040000/model_state/ocdbt.process_0/d/7dab326fb2cb1f0e157e155c7c8b04cc
new file mode 100644
index 0000000000000000000000000000000000000000..71a3c89a51b2461c4933e3feacd73658f93958a3
--- /dev/null
+++ b/040000/model_state/ocdbt.process_0/d/7dab326fb2cb1f0e157e155c7c8b04cc
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:159a712a86892f2ec32f915a97a5a32171df94fa62e008555e2b5f6aade743b2
+size 321597440
diff --git a/040000/model_state/ocdbt.process_0/d/de047e9eccf6eefb791a571eae771116 b/040000/model_state/ocdbt.process_0/d/de047e9eccf6eefb791a571eae771116
new file mode 100644
index 0000000000000000000000000000000000000000..39314fd88869afd0c395c47d9418a717243dffb6
--- /dev/null
+++ b/040000/model_state/ocdbt.process_0/d/de047e9eccf6eefb791a571eae771116
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:731fb7e9e3fa05695f609f673d2fd364704a97607eddd88c9c4d9791d431ccd5
+size 107578
diff --git a/060000/model_state/ocdbt.process_0/d/4ec4cbcd896ac4179022886603c5c967 b/060000/model_state/ocdbt.process_0/d/4ec4cbcd896ac4179022886603c5c967
new file mode 100644
index 0000000000000000000000000000000000000000..7fabefbced8496f5dfec61bfba5bec854d2933b7
--- /dev/null
+++ b/060000/model_state/ocdbt.process_0/d/4ec4cbcd896ac4179022886603c5c967
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:501a327b72420df507cbdda30121b3b244f757ab09d4261da7fae62fa36bd13d
+size 1047021
diff --git a/060000/model_state/ocdbt.process_0/d/5d6955d327966b061974e60856f65de2 b/060000/model_state/ocdbt.process_0/d/5d6955d327966b061974e60856f65de2
new file mode 100644
index 0000000000000000000000000000000000000000..c7a2871c7bd76521dc1d8866a2e8e6c014fde845
--- /dev/null
+++ b/060000/model_state/ocdbt.process_0/d/5d6955d327966b061974e60856f65de2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b95b5272a75ea725a1bed4a52a3e459aee580d07a041c4b9ce500a707f14ba2
+size 162848768
diff --git a/060000/model_state/ocdbt.process_0/d/c81fa37784c7b44640b76c45bd180d65 b/060000/model_state/ocdbt.process_0/d/c81fa37784c7b44640b76c45bd180d65
new file mode 100644
index 0000000000000000000000000000000000000000..3f269d0b61697c70fbe3eec82b8a6a66d66ef93f
--- /dev/null
+++ b/060000/model_state/ocdbt.process_0/d/c81fa37784c7b44640b76c45bd180d65
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b56fe5a79cd12e235d43a5251e1f786a272df5271adc1c2a69a915f2b78aaebf
+size 444170240
diff --git a/080000/model_state/ocdbt.process_0/d/2b34c992eeed9ced9ff9fc71c9299a07 b/080000/model_state/ocdbt.process_0/d/2b34c992eeed9ced9ff9fc71c9299a07
new file mode 100644
index 0000000000000000000000000000000000000000..0780430281f8bb8d74a9f688c1b5374b77330de1
--- /dev/null
+++ b/080000/model_state/ocdbt.process_0/d/2b34c992eeed9ced9ff9fc71c9299a07
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:566befabfc6bcc18ebdb646985cb5db96e03b82c3ccc99dc26d08a591574337a
+size 144642048
diff --git a/080000/model_state/ocdbt.process_0/d/709cd657dccf2c5adc568db343f42725 b/080000/model_state/ocdbt.process_0/d/709cd657dccf2c5adc568db343f42725
new file mode 100644
index 0000000000000000000000000000000000000000..7268a8582f610bd3e63c014427d9367013bf3ff3
--- /dev/null
+++ b/080000/model_state/ocdbt.process_0/d/709cd657dccf2c5adc568db343f42725
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be069e9c745fdeaf30a0482eb82df4a32732896803328450774d859bb31fedc9
+size 5365760
diff --git a/080000/model_state/ocdbt.process_0/d/7ea6d160856219e72ad3ac6bcd824739 b/080000/model_state/ocdbt.process_0/d/7ea6d160856219e72ad3ac6bcd824739
new file mode 100644
index 0000000000000000000000000000000000000000..a6340c0cadde672bc3105976ba74a9f14e96651b
--- /dev/null
+++ b/080000/model_state/ocdbt.process_0/d/7ea6d160856219e72ad3ac6bcd824739
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82a582d8f72ea1fc785a89fd0953374732f67d31bf9768b87d6668079df279b5
+size 457744384
diff --git a/100000/model_state/ocdbt.process_0/d/4908b67a51e8f6b19d57e06dbcc3a6d6 b/100000/model_state/ocdbt.process_0/d/4908b67a51e8f6b19d57e06dbcc3a6d6
new file mode 100644
index 0000000000000000000000000000000000000000..5fdaa640d89725b5617364c4430acf84a9c9e440
--- /dev/null
+++ b/100000/model_state/ocdbt.process_0/d/4908b67a51e8f6b19d57e06dbcc3a6d6
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ef10d04e03a9ead6f56b89dc848336370ac6fca17b733596e5bfc54a2b73dd5
+size 66043904
diff --git a/100000/model_state/ocdbt.process_0/d/98105609c0838c418037c239d750afd7 b/100000/model_state/ocdbt.process_0/d/98105609c0838c418037c239d750afd7
new file mode 100644
index 0000000000000000000000000000000000000000..30557e92a14c90c73fbe83e2ba9443da13378e56
--- /dev/null
+++ b/100000/model_state/ocdbt.process_0/d/98105609c0838c418037c239d750afd7
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe2913ce79bf5e41b4ea35bbac91f866e2a2bd7b4cea784e92d2dcd5bafb4f25
+size 541761536
diff --git a/120000/_CHECKPOINT_METADATA b/120000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e32704b366b2b17eea9ef36da8c172f860ebc530
--- /dev/null
+++ b/120000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"model_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler", "val_dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1761762052574653153, "commit_timestamp_nsecs": 1761762057517394585, "custom_metadata": {}}
\ No newline at end of file
diff --git a/120000/model_state/_METADATA b/120000/model_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..a530fd9c45497b1ccb6784d3056f6f309dcaa3d7
--- /dev/null
+++ b/120000/model_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('model', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('model', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('model', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('model', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('opt_state', '0', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'mu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'nu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '2', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step', 'value')": {"key_metadata": [{"key": "step", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/120000/model_state/_sharding b/120000/model_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..82fb9536f0260fad0437efaa6d41bd1f0ce05652
--- /dev/null
+++ b/120000/model_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjAuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjIuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","c3RlcC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/120000/model_state/array_metadatas/process_0 b/120000/model_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..d343705cd4e8951a39e35d6ea9e4bb34f47f7886
--- /dev/null
+++ b/120000/model_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "model.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "opt_state.0.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.2.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/120000/model_state/manifest.ocdbt b/120000/model_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..9f7acad10cde9acd75e8fdd0eaa1e383048f25c5
Binary files /dev/null and b/120000/model_state/manifest.ocdbt differ
diff --git a/120000/model_state/ocdbt.process_0/d/c4873dcbc98846803bc300a0644484d3 b/120000/model_state/ocdbt.process_0/d/c4873dcbc98846803bc300a0644484d3
new file mode 100644
index 0000000000000000000000000000000000000000..d8999c2c62c1ae6cbfe78a7d1abbf255c462e56e
--- /dev/null
+++ b/120000/model_state/ocdbt.process_0/d/c4873dcbc98846803bc300a0644484d3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50fb30c701ea9380195f723f8066d9376383fa3529fd0ed6eba9297f51100938
+size 202780672
diff --git a/120000/model_state/ocdbt.process_0/d/dca42bda6ce575cf0e3facd6ccefef2e b/120000/model_state/ocdbt.process_0/d/dca42bda6ce575cf0e3facd6ccefef2e
new file mode 100644
index 0000000000000000000000000000000000000000..244177406c77f1c45cac354e23a8dd684888d6bc
--- /dev/null
+++ b/120000/model_state/ocdbt.process_0/d/dca42bda6ce575cf0e3facd6ccefef2e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:007ec922c341e2330a337d08e652aa796a45d2312653ec6def5b5469c504a5c9
+size 3959303
diff --git a/120000/model_state/ocdbt.process_0/d/fa3cb93f242c469029cb4005bed01132 b/120000/model_state/ocdbt.process_0/d/fa3cb93f242c469029cb4005bed01132
new file mode 100644
index 0000000000000000000000000000000000000000..40f8596e6e03a3e414f9f7e58d364b69d750a5ad
--- /dev/null
+++ b/120000/model_state/ocdbt.process_0/d/fa3cb93f242c469029cb4005bed01132
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe2462672510f6e4e2ecc8e8717f3937897eca9ef5238bf39b020a523d31ab75
+size 401580032
diff --git a/120000/model_state/ocdbt.process_0/manifest.ocdbt b/120000/model_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..1ba345943c00d21ea78251e07f04ca004a815690
Binary files /dev/null and b/120000/model_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/120000/train_dataloader_state/process_0-of-1.json b/120000/train_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..a0b61780ae83a61b8d12eecf0a4df303b52e38b7
--- /dev/null
+++ b/120000/train_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 4344656,
+        "1": 4345137,
+        "2": 4345290,
+        "3": 4345611,
+        "4": 4345284,
+        "5": 4345597,
+        "6": 4344854,
+        "7": 4344887
+    },
+    "last_worker_index": 7,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=613100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=b4eb1f2c7517a10c44280d95cdea999338e38889)"
+}
\ No newline at end of file
diff --git a/120000/val_dataloader_state/process_0-of-1.json b/120000/val_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..c01929ddc725a2a2bf026fec73283cbfb15da318
--- /dev/null
+++ b/120000/val_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 11320,
+        "1": 11289,
+        "2": 10986,
+        "3": 10995,
+        "4": 11020,
+        "5": 11021,
+        "6": 10982,
+        "7": 11031
+    },
+    "last_worker_index": 1,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=6100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=c7f9721eaa96e2b50f0ce20aaf09ae05f4f0d146)"
+}
\ No newline at end of file
diff --git a/140000/model_state/_METADATA b/140000/model_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..a530fd9c45497b1ccb6784d3056f6f309dcaa3d7
--- /dev/null
+++ b/140000/model_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('model', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('model', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('model', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('model', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('opt_state', '0', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'mu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'nu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '2', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step', 'value')": {"key_metadata": [{"key": "step", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/140000/model_state/_sharding b/140000/model_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..82fb9536f0260fad0437efaa6d41bd1f0ce05652
--- /dev/null
+++ b/140000/model_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjAuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjIuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","c3RlcC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/140000/model_state/d/8a276b19494cdd7a9a73ecd3581b7428 b/140000/model_state/d/8a276b19494cdd7a9a73ecd3581b7428
new file mode 100644
index 0000000000000000000000000000000000000000..e6c72d568dcab065322ec64495cb6e291b318e79
Binary files /dev/null and b/140000/model_state/d/8a276b19494cdd7a9a73ecd3581b7428 differ
diff --git a/140000/model_state/ocdbt.process_0/d/0a158a7ea66c6d711deea52de8ce0a32 b/140000/model_state/ocdbt.process_0/d/0a158a7ea66c6d711deea52de8ce0a32
new file mode 100644
index 0000000000000000000000000000000000000000..cc6362f8dda9c0e28513965f1c9cecdff0256d35
Binary files /dev/null and b/140000/model_state/ocdbt.process_0/d/0a158a7ea66c6d711deea52de8ce0a32 differ
diff --git a/140000/model_state/ocdbt.process_0/d/23085a1042588716041c3d37e0492fed b/140000/model_state/ocdbt.process_0/d/23085a1042588716041c3d37e0492fed
new file mode 100644
index 0000000000000000000000000000000000000000..0cae882758bfe9b59ec29ad83739436357f0e58e
--- /dev/null
+++ b/140000/model_state/ocdbt.process_0/d/23085a1042588716041c3d37e0492fed
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7692cb37fd3412d57981254f1f6bb92f92fe0760f927fc48ea6a09952f64ebda
+size 89255936
diff --git a/140000/model_state/ocdbt.process_0/d/419686858425dc7d2bca5735a2781a05 b/140000/model_state/ocdbt.process_0/d/419686858425dc7d2bca5735a2781a05
new file mode 100644
index 0000000000000000000000000000000000000000..939393f423f5983ae5d72551c1e472ef918efa95
Binary files /dev/null and b/140000/model_state/ocdbt.process_0/d/419686858425dc7d2bca5735a2781a05 differ
diff --git a/140000/model_state/ocdbt.process_0/d/57dff85bed1318fc6766dc75f0bbf7f2 b/140000/model_state/ocdbt.process_0/d/57dff85bed1318fc6766dc75f0bbf7f2
new file mode 100644
index 0000000000000000000000000000000000000000..cbb96ff11a49c084d348ce46d42d0f3153367478
Binary files /dev/null and b/140000/model_state/ocdbt.process_0/d/57dff85bed1318fc6766dc75f0bbf7f2 differ
diff --git a/140000/model_state/ocdbt.process_0/d/8dd46fb95b571ed876e66163c846288e b/140000/model_state/ocdbt.process_0/d/8dd46fb95b571ed876e66163c846288e
new file mode 100644
index 0000000000000000000000000000000000000000..e0e384c5451af3f8eb59f03a594dfebc0af9c877
--- /dev/null
+++ b/140000/model_state/ocdbt.process_0/d/8dd46fb95b571ed876e66163c846288e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b45e66fe51e574dfedce23549438548e92fc1bfa4e566c062a0610a82a417079
+size 30154752
diff --git a/140000/model_state/ocdbt.process_0/d/975df04c2b623aebb500442c29de2a84 b/140000/model_state/ocdbt.process_0/d/975df04c2b623aebb500442c29de2a84
new file mode 100644
index 0000000000000000000000000000000000000000..154cc9e6f8c6088ea201d585007e44868fd3e9bc
Binary files /dev/null and b/140000/model_state/ocdbt.process_0/d/975df04c2b623aebb500442c29de2a84 differ
diff --git a/140000/model_state/ocdbt.process_0/d/9974f401c505e80ccb31e492612019d9 b/140000/model_state/ocdbt.process_0/d/9974f401c505e80ccb31e492612019d9
new file mode 100644
index 0000000000000000000000000000000000000000..85871beeb40d7d7538bd3d4e3ac20c05964a9dd2
--- /dev/null
+++ b/140000/model_state/ocdbt.process_0/d/9974f401c505e80ccb31e492612019d9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21c0ae6a46cc8b0ff9d2a91ff0cdf10facbf1031d6a6bf05d395718f32254271
+size 488345600
diff --git a/140000/model_state/ocdbt.process_0/d/b5056db105a7ee9a2d5c0d8407627db5 b/140000/model_state/ocdbt.process_0/d/b5056db105a7ee9a2d5c0d8407627db5
new file mode 100644
index 0000000000000000000000000000000000000000..308a2468ae9d7475962caaf00f7b213d4f830d4a
Binary files /dev/null and b/140000/model_state/ocdbt.process_0/d/b5056db105a7ee9a2d5c0d8407627db5 differ
diff --git a/140000/val_dataloader_state/process_0-of-1.json b/140000/val_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..5a0ce0e1454c2db3f4c179b178de70dfc1588b6d
--- /dev/null
+++ b/140000/val_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 13056,
+        "1": 13025,
+        "2": 13010,
+        "3": 13027,
+        "4": 13044,
+        "5": 12765,
+        "6": 12718,
+        "7": 12783
+    },
+    "last_worker_index": 4,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=6100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=c7f9721eaa96e2b50f0ce20aaf09ae05f4f0d146)"
+}
\ No newline at end of file
diff --git a/160000/_CHECKPOINT_METADATA b/160000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ab3ab15a11abd5b1a227f86096b5f735f41531a9
--- /dev/null
+++ b/160000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"model_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler", "val_dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1761766237445022179, "commit_timestamp_nsecs": 1761766242903626218, "custom_metadata": {}}
\ No newline at end of file
diff --git a/160000/model_state/_METADATA b/160000/model_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..a530fd9c45497b1ccb6784d3056f6f309dcaa3d7
--- /dev/null
+++ b/160000/model_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('model', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('model', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('model', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('model', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('opt_state', '0', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'mu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'nu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '2', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step', 'value')": {"key_metadata": [{"key": "step", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/160000/model_state/_sharding b/160000/model_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..82fb9536f0260fad0437efaa6d41bd1f0ce05652
--- /dev/null
+++ b/160000/model_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjAuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjIuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","c3RlcC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/160000/model_state/array_metadatas/process_0 b/160000/model_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..d343705cd4e8951a39e35d6ea9e4bb34f47f7886
--- /dev/null
+++ b/160000/model_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "model.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "opt_state.0.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.2.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/160000/model_state/d/217a9f1e54011c52a58264fea1140eda b/160000/model_state/d/217a9f1e54011c52a58264fea1140eda
new file mode 100644
index 0000000000000000000000000000000000000000..62169d52224d09a44c58a50a58820db6401a7d2a
Binary files /dev/null and b/160000/model_state/d/217a9f1e54011c52a58264fea1140eda differ
diff --git a/160000/model_state/manifest.ocdbt b/160000/model_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..a8df3d700d3951bca229feb386c846833b6e24ab
Binary files /dev/null and b/160000/model_state/manifest.ocdbt differ
diff --git a/160000/model_state/ocdbt.process_0/d/179d8e0ef0487f376bbe06c6778e9a40 b/160000/model_state/ocdbt.process_0/d/179d8e0ef0487f376bbe06c6778e9a40
new file mode 100644
index 0000000000000000000000000000000000000000..2494ae7f0883be8a343c01695c751affc72878c9
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/d/179d8e0ef0487f376bbe06c6778e9a40 differ
diff --git a/160000/model_state/ocdbt.process_0/d/1cfd62eb21313e54e9e6f1f01d1e15de b/160000/model_state/ocdbt.process_0/d/1cfd62eb21313e54e9e6f1f01d1e15de
new file mode 100644
index 0000000000000000000000000000000000000000..fda624d68d26febe850104cc3038abcf6eaa8402
--- /dev/null
+++ b/160000/model_state/ocdbt.process_0/d/1cfd62eb21313e54e9e6f1f01d1e15de
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c43644c4d9358cf9970f250b8bb6a1594622301a367e94cfbd0b441b47a5019
+size 47153152
diff --git a/160000/model_state/ocdbt.process_0/d/2b6bc3528951c6b5c3804eea00f4ed09 b/160000/model_state/ocdbt.process_0/d/2b6bc3528951c6b5c3804eea00f4ed09
new file mode 100644
index 0000000000000000000000000000000000000000..308a2468ae9d7475962caaf00f7b213d4f830d4a
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/d/2b6bc3528951c6b5c3804eea00f4ed09 differ
diff --git a/160000/model_state/ocdbt.process_0/d/755249cb63c86652a08a74a9e414216e b/160000/model_state/ocdbt.process_0/d/755249cb63c86652a08a74a9e414216e
new file mode 100644
index 0000000000000000000000000000000000000000..cc6362f8dda9c0e28513965f1c9cecdff0256d35
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/d/755249cb63c86652a08a74a9e414216e differ
diff --git a/160000/model_state/ocdbt.process_0/d/9aef4849f854df2df779cfb1c0555793 b/160000/model_state/ocdbt.process_0/d/9aef4849f854df2df779cfb1c0555793
new file mode 100644
index 0000000000000000000000000000000000000000..46e7b63a695a97a6eaeae945f63af18147d95681
--- /dev/null
+++ b/160000/model_state/ocdbt.process_0/d/9aef4849f854df2df779cfb1c0555793
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:290209c3dd4104eb8dc21aa2c5bb5a4a1a116426c43f3d1f675a8da0c9607833
+size 519970816
diff --git a/160000/model_state/ocdbt.process_0/d/9e1ef595d4e11b2fc1a9c26f5880e9cf b/160000/model_state/ocdbt.process_0/d/9e1ef595d4e11b2fc1a9c26f5880e9cf
new file mode 100644
index 0000000000000000000000000000000000000000..d015affdbc6a14ba215d8f663ce360d24978bb34
--- /dev/null
+++ b/160000/model_state/ocdbt.process_0/d/9e1ef595d4e11b2fc1a9c26f5880e9cf
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79d0ce6d1737c74aed6ee05947f26ef8a50b23a041850313e434afe5074bd99f
+size 34938880
diff --git a/160000/model_state/ocdbt.process_0/d/c99ef08bc07f9034eae0283f1860048a b/160000/model_state/ocdbt.process_0/d/c99ef08bc07f9034eae0283f1860048a
new file mode 100644
index 0000000000000000000000000000000000000000..e53bbc46f435f7f7c9b0a291465c170e4c73967c
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/d/c99ef08bc07f9034eae0283f1860048a differ
diff --git a/160000/model_state/ocdbt.process_0/d/ce5c874b15e827392c30c984a16dce3a b/160000/model_state/ocdbt.process_0/d/ce5c874b15e827392c30c984a16dce3a
new file mode 100644
index 0000000000000000000000000000000000000000..79395e76405e421f389628f8025fe67057cf1eab
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/d/ce5c874b15e827392c30c984a16dce3a differ
diff --git a/160000/model_state/ocdbt.process_0/d/e607017ce17c1af8daad4f13214b3ea6 b/160000/model_state/ocdbt.process_0/d/e607017ce17c1af8daad4f13214b3ea6
new file mode 100644
index 0000000000000000000000000000000000000000..b8c5a74da598dc66e1e13c71051aa8d51e30ff8d
--- /dev/null
+++ b/160000/model_state/ocdbt.process_0/d/e607017ce17c1af8daad4f13214b3ea6
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbd7ce7e1a8b31a01e685e48d242b78449fa03fd27af24e908a084c5605bcb08
+size 5914624
diff --git a/160000/model_state/ocdbt.process_0/manifest.ocdbt b/160000/model_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..79c8d01a99b35063ca6204ec7752afc0d298c73b
Binary files /dev/null and b/160000/model_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/160000/train_dataloader_state/process_0-of-1.json b/160000/train_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..f434fefbc8f12792f064d8a486882b140da1fd51
--- /dev/null
+++ b/160000/train_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 5792608,
+        "1": 5793521,
+        "2": 5793738,
+        "3": 5794603,
+        "4": 5793732,
+        "5": 5794237,
+        "6": 5792622,
+        "7": 5793127
+    },
+    "last_worker_index": 7,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=613100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=b4eb1f2c7517a10c44280d95cdea999338e38889)"
+}
\ No newline at end of file
diff --git a/160000/val_dataloader_state/process_0-of-1.json b/160000/val_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..09c2a19bf6fea42feaa59b04226fa6bfdb8e3670
--- /dev/null
+++ b/160000/val_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 14784,
+        "1": 14809,
+        "2": 14754,
+        "3": 14755,
+        "4": 14788,
+        "5": 14797,
+        "6": 14742,
+        "7": 14815
+    },
+    "last_worker_index": 7,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=6100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=c7f9721eaa96e2b50f0ce20aaf09ae05f4f0d146)"
+}
\ No newline at end of file
diff --git a/180000/model_state/ocdbt.process_0/d/b8a20e22aa870a02f9ebdd120566cee1 b/180000/model_state/ocdbt.process_0/d/b8a20e22aa870a02f9ebdd120566cee1
new file mode 100644
index 0000000000000000000000000000000000000000..3f76c4a8518d5184de286886f59659a501bea480
--- /dev/null
+++ b/180000/model_state/ocdbt.process_0/d/b8a20e22aa870a02f9ebdd120566cee1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4868ebff2cc64201e97311f4085df98ab7d3579834df8227ec8faf88086dee75
+size 71045120
diff --git a/180000/model_state/ocdbt.process_0/d/ffe4cb9a3ef849cb1dd1ad05694ceb3f b/180000/model_state/ocdbt.process_0/d/ffe4cb9a3ef849cb1dd1ad05694ceb3f
new file mode 100644
index 0000000000000000000000000000000000000000..027d9dcd385cc38816d693ff4ddfafd1c4c5ff6e
--- /dev/null
+++ b/180000/model_state/ocdbt.process_0/d/ffe4cb9a3ef849cb1dd1ad05694ceb3f
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14c957cc1e28b4e58d5e0f5424c45e1e1df976f8c4e05d7004eecebd43b344ee
+size 537108480
diff --git a/190000/model_state/ocdbt.process_0/d/682083c1dd6cfab1449e6f9c0956b612 b/190000/model_state/ocdbt.process_0/d/682083c1dd6cfab1449e6f9c0956b612
new file mode 100644
index 0000000000000000000000000000000000000000..84d048188a4476ff53be038aa410d5ff7a39a40a
--- /dev/null
+++ b/190000/model_state/ocdbt.process_0/d/682083c1dd6cfab1449e6f9c0956b612
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb2ba1de6d962bbe94b556983869a1c9a77cda41d4e94187a2dfbec9a34800cb
+size 1068276
diff --git a/190000/model_state/ocdbt.process_0/d/90986d89119a1a9d0a93a58c8ffd3aaf b/190000/model_state/ocdbt.process_0/d/90986d89119a1a9d0a93a58c8ffd3aaf
new file mode 100644
index 0000000000000000000000000000000000000000..2759ca60722fbd8967a47f68b5d6a8a52dfe8687
--- /dev/null
+++ b/190000/model_state/ocdbt.process_0/d/90986d89119a1a9d0a93a58c8ffd3aaf
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd3d3dd059cb1e6f6e34788acb379798e69372827771cd472cf5a11abcc83316
+size 29114368
diff --git a/190000/model_state/ocdbt.process_0/d/aba66df9f95c5f517e76eec83302d73d b/190000/model_state/ocdbt.process_0/d/aba66df9f95c5f517e76eec83302d73d
new file mode 100644
index 0000000000000000000000000000000000000000..8be2d064c93d79b570ace4e14df68389cd041624
--- /dev/null
+++ b/190000/model_state/ocdbt.process_0/d/aba66df9f95c5f517e76eec83302d73d
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9866ddcfdffe78f860eff85f56b7bcfe353efd73b92639bcd7a7cb9b37fbdcee
+size 577675264
diff --git a/195000/model_state/ocdbt.process_0/d/805b4d2cbff076e49d4272cb4de1c3c2 b/195000/model_state/ocdbt.process_0/d/805b4d2cbff076e49d4272cb4de1c3c2
new file mode 100644
index 0000000000000000000000000000000000000000..058627ff3425ff2a8a448e6987309acca84ca81b
--- /dev/null
+++ b/195000/model_state/ocdbt.process_0/d/805b4d2cbff076e49d4272cb4de1c3c2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a5a480070c1c1bce32fe344c6a18a985823bcaff1e6d0cbb877cc5d11e074ad
+size 44761088
diff --git a/195000/model_state/ocdbt.process_0/d/8d25c17882bc09ef12ec3d5d753f6a4a b/195000/model_state/ocdbt.process_0/d/8d25c17882bc09ef12ec3d5d753f6a4a
new file mode 100644
index 0000000000000000000000000000000000000000..6b86ac5365c570b46a2c80593343170202f6dada
--- /dev/null
+++ b/195000/model_state/ocdbt.process_0/d/8d25c17882bc09ef12ec3d5d753f6a4a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eceb1fd3263a43b29fac84243ad46247008caf621e6dae1a949dabd99599cf93
+size 102875136
diff --git a/195000/model_state/ocdbt.process_0/d/9a90ed05f2a1c9ce06491436f3ef894c b/195000/model_state/ocdbt.process_0/d/9a90ed05f2a1c9ce06491436f3ef894c
new file mode 100644
index 0000000000000000000000000000000000000000..7b66bcca74063a2fb7a648462a3bfcd9a7178714
--- /dev/null
+++ b/195000/model_state/ocdbt.process_0/d/9a90ed05f2a1c9ce06491436f3ef894c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b17bd7bcf60ea20da34a89f2b97edccd3c0e6fbd1c06fadbda9b58162e6475df
+size 460075008
diff --git a/200000/model_state/_METADATA b/200000/model_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..a530fd9c45497b1ccb6784d3056f6f309dcaa3d7
--- /dev/null
+++ b/200000/model_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('model', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('model', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('model', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('model', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('model', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'lam', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('model', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('model', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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": []}}, "('model', 'tokenizer', 'vq', 'drop', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "drop", "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]}}, "('opt_state', '0', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'mu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'mu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'mu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'mask_token', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "mask_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'patch_embed', 'embedding', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "patch_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": [1024, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '4', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'blocks', '5', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_dense", "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]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '0', 'nu', 'dynamics', 'transformer', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('opt_state', '0', 'nu', 'lam', 'action_in', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_in", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 1, 768]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'action_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "action_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'lam', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'lam', 'patch_up', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "patch_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'lam', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "lam", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'decoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "output_dense", "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, 768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '0', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '1', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '2', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense1', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense1", "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, 2048]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_dense2', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_dense2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ffn_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'spatial_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "spatial_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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": [8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_attention', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_attention", "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, 8, 64]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'blocks', '3', 'temporal_norm', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "temporal_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]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm1', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'input_norm2', 'scale', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "input_norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'bias', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'encoder', 'output_dense', 'kernel', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "output_dense", "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, 32]}}, "('opt_state', '0', 'nu', 'tokenizer', 'vq', 'codebook', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "tokenizer", "key_type": 2}, {"key": "vq", "key_type": 2}, {"key": "codebook", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 32]}}, "('opt_state', '2', 'count', 'value')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step', 'value')": {"key_metadata": [{"key": "step", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/200000/model_state/_sharding b/200000/model_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..82fb9536f0260fad0437efaa6d41bd1f0ce05652
--- /dev/null
+++ b/200000/model_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjAuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubXUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjAubnUudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","b3B0X3N0YXRlLjIuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MuYWN0aW9uX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MubWFza190b2tlbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MucGF0Y2hfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjQudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuYmxvY2tzLjUudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwuZHluYW1pY3MudHJhbnNmb3JtZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl91cC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmFjdGlvbl9pbi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnBhdGNoX3VwLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwubGFtLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmRlY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjAudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjEudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjIudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuZmZuX2RlbnNlMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9hdHRlbnRpb24udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMuc3BhdGlhbF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfYXR0ZW50aW9uLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuYmxvY2tzLjMudGVtcG9yYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2UuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfZGVuc2Uua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIuaW5wdXRfbm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLmVuY29kZXIub3V0cHV0X2RlbnNlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmNvZGVib29rLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","bW9kZWwudG9rZW5pemVyLnZxLmRyb3Aucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","c3RlcC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/200000/model_state/array_metadatas/process_0 b/200000/model_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..d343705cd4e8951a39e35d6ea9e4bb34f47f7886
--- /dev/null
+++ b/200000/model_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "model.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "opt_state.0.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.mask_token.value", "write_shape": [1, 1, 1, 512], "chunk_shape": [1, 1, 1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.patch_embed.embedding.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_dense.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.transformer.output_dense.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.vq.codebook.value", "write_shape": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.2.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/200000/model_state/d/64125bb679de010ad204ad5d723fd3a7 b/200000/model_state/d/64125bb679de010ad204ad5d723fd3a7
new file mode 100644
index 0000000000000000000000000000000000000000..a0da1ea674276dae92b25a3203c5995d7e1ba39e
Binary files /dev/null and b/200000/model_state/d/64125bb679de010ad204ad5d723fd3a7 differ
diff --git a/200000/model_state/manifest.ocdbt b/200000/model_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..8e3baa704a04cc1c6f881cef9836110b1967d0dc
Binary files /dev/null and b/200000/model_state/manifest.ocdbt differ
diff --git a/200000/model_state/ocdbt.process_0/d/063057d29f7537fe5dc89d2c5438d296 b/200000/model_state/ocdbt.process_0/d/063057d29f7537fe5dc89d2c5438d296
new file mode 100644
index 0000000000000000000000000000000000000000..2791cf9909b010f0211d57cf75338fbaf200aa80
--- /dev/null
+++ b/200000/model_state/ocdbt.process_0/d/063057d29f7537fe5dc89d2c5438d296
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b211dd4907daf99529f201337a116c7929379c9767a7191a70172b86c30ee6ea
+size 53641216
diff --git a/200000/model_state/ocdbt.process_0/d/2beb4872931d211d1a884a7b30f6ed44 b/200000/model_state/ocdbt.process_0/d/2beb4872931d211d1a884a7b30f6ed44
new file mode 100644
index 0000000000000000000000000000000000000000..b4b55d7adc82f8fe8add582449d23ced7de37aff
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/d/2beb4872931d211d1a884a7b30f6ed44 differ
diff --git a/200000/model_state/ocdbt.process_0/d/45b803694d1df8de743a6714a80a52e5 b/200000/model_state/ocdbt.process_0/d/45b803694d1df8de743a6714a80a52e5
new file mode 100644
index 0000000000000000000000000000000000000000..3437aa20a62ded06a952eed648ef7a3f0a8e56f6
--- /dev/null
+++ b/200000/model_state/ocdbt.process_0/d/45b803694d1df8de743a6714a80a52e5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c061bab29ada458df795c9ae39b94ff35b0a768e79e8607df631a0a44b5f694
+size 980907
diff --git a/200000/model_state/ocdbt.process_0/d/4af41dbb6e42fbb754c7f0c8a5199041 b/200000/model_state/ocdbt.process_0/d/4af41dbb6e42fbb754c7f0c8a5199041
new file mode 100644
index 0000000000000000000000000000000000000000..308a2468ae9d7475962caaf00f7b213d4f830d4a
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/d/4af41dbb6e42fbb754c7f0c8a5199041 differ
diff --git a/200000/model_state/ocdbt.process_0/d/9b89c207d336038085afc24e415e4801 b/200000/model_state/ocdbt.process_0/d/9b89c207d336038085afc24e415e4801
new file mode 100644
index 0000000000000000000000000000000000000000..0bab7cdb6dad39f4e4f0127ca421fcf4707af637
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/d/9b89c207d336038085afc24e415e4801 differ
diff --git a/200000/model_state/ocdbt.process_0/d/d9cc1eb28602948dbae8c69e98e97806 b/200000/model_state/ocdbt.process_0/d/d9cc1eb28602948dbae8c69e98e97806
new file mode 100644
index 0000000000000000000000000000000000000000..cc6362f8dda9c0e28513965f1c9cecdff0256d35
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/d/d9cc1eb28602948dbae8c69e98e97806 differ
diff --git a/200000/model_state/ocdbt.process_0/d/de85c51af409fc8267cf059f45307e01 b/200000/model_state/ocdbt.process_0/d/de85c51af409fc8267cf059f45307e01
new file mode 100644
index 0000000000000000000000000000000000000000..0eb5be944ea01a67719c8457722759c64ac5c851
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/d/de85c51af409fc8267cf059f45307e01 differ
diff --git a/200000/model_state/ocdbt.process_0/d/ee516710be0ffac27c9c58605cf19264 b/200000/model_state/ocdbt.process_0/d/ee516710be0ffac27c9c58605cf19264
new file mode 100644
index 0000000000000000000000000000000000000000..a12e5786490794d4b43745fa5b4978d92ed41a31
--- /dev/null
+++ b/200000/model_state/ocdbt.process_0/d/ee516710be0ffac27c9c58605cf19264
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91cd6ff1ff40c550e83512a6ab3fe3110c342342186ecd722faf378948fbdb6e
+size 553316352
diff --git a/200000/model_state/ocdbt.process_0/manifest.ocdbt b/200000/model_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..80dd7b041123ad88ce306167cb1a346d6fdf511d
Binary files /dev/null and b/200000/model_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/200000/train_dataloader_state/process_0-of-1.json b/200000/train_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..936ee185416f55df1c4e2280d6690ffed48978f3
--- /dev/null
+++ b/200000/train_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 7241008,
+        "1": 7242097,
+        "2": 7241978,
+        "3": 7242883,
+        "4": 7242212,
+        "5": 7242485,
+        "6": 7241326,
+        "7": 7241471
+    },
+    "last_worker_index": 7,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=613100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=b4eb1f2c7517a10c44280d95cdea999338e38889)"
+}
\ No newline at end of file
diff --git a/200000/val_dataloader_state/process_0-of-1.json b/200000/val_dataloader_state/process_0-of-1.json
new file mode 100644
index 0000000000000000000000000000000000000000..68fd46331bd1eefdbc2b7f8559ebb294f19639a7
--- /dev/null
+++ b/200000/val_dataloader_state/process_0-of-1.json
@@ -0,0 +1,17 @@
+{
+    "version": 2,
+    "last_seen_indices": {
+        "0": 18536,
+        "1": 18569,
+        "2": 18522,
+        "3": 18515,
+        "4": 18548,
+        "5": 18573,
+        "6": 18206,
+        "7": 18303
+    },
+    "last_worker_index": 5,
+    "worker_count": 8,
+    "sampler": "IndexSampler(num_records=6100, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
+    "data_source": "ArrayRecordDataSource(hash_of_paths=c7f9721eaa96e2b50f0ce20aaf09ae05f4f0d146)"
+}
\ No newline at end of file